package com.gutsyzhan.yusiailangchain4j.controller;

import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.splitter.DocumentByLineSplitter;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.onnx.HuggingFaceTokenizer;
import dev.langchain4j.model.ollama.OllamaEmbeddingModel;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.store.embedding.*;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

@Slf4j
@RestController
@RequestMapping("/embedding")
public class EmbeddingController {

    @Autowired
    private OllamaEmbeddingModel ollamaEmbeddingModel;

    @Autowired
    private EmbeddingStore embeddingStore;

    @GetMapping("/readAndStore")
    public String readAndStore(){
        String filePath = "E:/Files/langchain4j/files/langchain4j.txt";
        // 加载文档
        Document document = FileSystemDocumentLoader.loadDocument(filePath);
        // 创建内存向量存储
        InMemoryEmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore();
        // 文档分割与嵌入生成
        EmbeddingStoreIngestor.ingest(document, embeddingStore);
        // 打印查看向量存储的内容
        log.info("EmbeddingStore: {}", embeddingStore);
        return embeddingStore.toString();
    }

    @GetMapping("/splitterAndStore")
    public String splitterAndStore(){
        String filePath = "E:/Files/langchain4j/files/langchain4j.txt";
        // 加载文档
        Document document = FileSystemDocumentLoader.loadDocument(filePath);
        // 创建内存向量存储
        InMemoryEmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore();
        // 创建一个文本分块器，将文本分块为多个段落
        DocumentByLineSplitter splitter = new DocumentByLineSplitter(300, 30);
        // 文档分割与嵌入生成
        EmbeddingStoreIngestor.builder()
                        .embeddingStore(embeddingStore)
                                .documentSplitter(splitter)
                                        .build()
                                              .ingest(document);
        // 打印查看向量存储的内容
        log.info("EmbeddingStore: {}", embeddingStore);
        return embeddingStore.toString();
    }

    @GetMapping("/tokenCount")
    public String tokenCount(){
        String text = "这是一个示例文本，用于测试 token 长度的计算。";
        // 创建用户消息
        UserMessage message = UserMessage.userMessage(text);
        // 创建一个 HuggingFaceTokenizer 对象
        HuggingFaceTokenizer  tokenizer = new HuggingFaceTokenizer();
        // 计算文本的 token 长度
        int tokenCount = tokenizer.estimateTokenCountInMessage(message);
        return tokenCount + "";
    }

    /**
     * 测试向量模型
     */
    @GetMapping("/testEmbeddingModel")
    public String testEmbeddingModel(){
        String text = "你好";
        Response<Embedding> response = ollamaEmbeddingModel.embed(text);
        log.info("向量：{}",response.content().vector().length);
        log.info("向量维度：{}",response.content().dimension());
        log.info("token用量：{}",response.tokenUsage());
        log.info("finishReason：{}",response.finishReason());
        log.info("向量输出：{}" ,response);
        return response.toString();
    }

    /**
     * 测试向量存储
     */
    @GetMapping("/testEmbeddingStore")
    public String testEmbeddingStore(){
        //将文本转换成向量
        TextSegment textSegment = TextSegment.from("我喜欢西瓜");
        //获取向量
        Embedding embedding = ollamaEmbeddingModel.embed(textSegment).content();
        //添加向量
        embeddingStore.add(embedding, textSegment);
        return "success";
    }

    /**
     * 测试向量搜索
     */
    @GetMapping("/testEmbeddingSearch")
    public String testEmbeddingSearch(){
        //提问，并将问题转成向量数据
        Embedding queryEmbedding = ollamaEmbeddingModel.embed("你最喜欢什么水果？").content();
        //创建搜索请求对象
        EmbeddingSearchRequest searchRequest = EmbeddingSearchRequest.builder()
                .queryEmbedding(queryEmbedding)
                .maxResults(1) // 设置最大返回结果数量
                .minScore(0.1) // 设置最小得分阈值
                .build();
        //根据搜索请求 searchRequest 在向量存储中进行相似度搜索
        EmbeddingSearchResult<TextSegment> searchResult = embeddingStore.search(searchRequest);
        //获取搜索结果,获取最相似的向量,返回一个
        EmbeddingMatch<TextSegment> embeddingMatch = searchResult.matches().get(0);
        //获取匹配项的相似度得分
        log.info("相似度得分：{}",embeddingMatch.score());
        // 获取匹配项的向量数据
        String text = embeddingMatch.embedded().text();
        //获取匹配项的元数据,即文本结果
        log.info("匹配项的元数据：{}",text);
        return text;
    }
}
